package com.agentai.framework.websocket;

import com.alibaba.fastjson.JSONObject;
import com.agentai.framework.common.utils.IpUtils;
import com.agentai.framework.common.utils.Result;
import com.agentai.framework.satoken.cache.TokenStoreCache;
import com.agentai.framework.security.user.UserDetail;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.HttpHeaders;
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
import io.netty.util.AttributeKey;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/**
 * @Author yangfei
 * @Description
 * @Date 2024/5/21 10:34
 */
@Component
@ChannelHandler.Sharable
@Slf4j
public class WebSocketHandler extends SimpleChannelInboundHandler<TextWebSocketFrame> {

    /**
     * websocket的token 别名
     */
    private static final String TOKEN_KEY = "Sec-WebSocket-Protocol";
    /**
     * websocket的 client_id 别名
     */
    private static final String WS_CLIENT = "WebSocket-client-id";
    @Autowired
    WebSocketConfig webSocketConfig;
    @Autowired
    TokenStoreCache tokenStoreCache;

    /**
     * 一旦连接，第一个被执行
     */
    @Override
    public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
        log.info("有新的客户端链接：[{}]", ctx.channel().id().asLongText());
        // 添加到channelGroup 通道组
        NettyConfig.getChannelGroup().add(ctx.channel());
    }

    /**
     * 读取数据
     */
    @Override
    protected void channelRead0(ChannelHandlerContext ctx, TextWebSocketFrame msg) throws Exception {

        log.info("服务器收到消息：{}", msg.text());

        if (msg.text().equals("ping")) {
            //心跳消息
            ctx.channel().writeAndFlush(new TextWebSocketFrame(JSONObject.toJSONString(Result.ok(null, "心跳响应"))));
        } else {
            // 回复消息
            ctx.channel().writeAndFlush(new TextWebSocketFrame(JSONObject.toJSONString(Result.ok())));
        }
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
        // 检测事件，如果是握手成功事件，做点业务处理
        if (evt instanceof WebSocketServerProtocolHandler.HandshakeComplete) {

            WebSocketServerProtocolHandler.HandshakeComplete complete = (WebSocketServerProtocolHandler.HandshakeComplete) evt;
            String uri = complete.requestUri(); // 获取请求路径，接着根据自己的业务逻辑判断以及获取自己想要的参数
            String regex = webSocketConfig.getPath();
            if (!(uri.startsWith(regex) && uri.contains(regex))) {
                log.info("请求路径不合法：{}", uri);
                ctx.channel().writeAndFlush(new TextWebSocketFrame(JSONObject.toJSONString(Result.error("请求路径不合法"))));
                ctx.close();
                return;
            }
            HttpHeaders headers = complete.requestHeaders();
            String userId = headers.get(WS_CLIENT);
            boolean isAuth = webSocketConfig.isAuthEnabled();
            if (isAuth) {
                String token = headers.get(TOKEN_KEY);
                UserDetail user = tokenStoreCache.getUser(token);
                if (null == user) {
                    log.info("Websocket连接 Token安全校验失败，IP:{}, Token:{}, Path = {}", IpUtils.getHostIp(), token, uri);
                    ctx.channel().writeAndFlush(new TextWebSocketFrame(JSONObject.toJSONString(Result.error("Token安全校验失败"))));
                    ctx.close();
                    return;
                }
                userId = String.valueOf(user.getId());
            }
            NettyConfig.getChannelMap().put(userId, ctx.channel());
            // 将用户ID作为自定义属性加入到channel中，方便随时channel中获取用户ID
            AttributeKey<String> key = AttributeKey.valueOf("userId");
            ctx.channel().attr(key).setIfAbsent(userId);
        } else {
            super.userEventTriggered(ctx, evt);
        }
    }


    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
        log.info("用户下线了:{}", ctx.channel().id().asLongText());
        // 删除通道
        NettyConfig.getChannelGroup().remove(ctx.channel());
        removeUserId(ctx);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        log.info("异常：{}", cause.getMessage());
        // 删除通道
        NettyConfig.getChannelGroup().remove(ctx.channel());
        removeUserId(ctx);
        ctx.close();
    }

    /**
     * 删除用户与channel的对应关系
     */
    private void removeUserId(ChannelHandlerContext ctx) {
        AttributeKey<String> key = AttributeKey.valueOf("userId");
        String userId = ctx.channel().attr(key).get();
        NettyConfig.getChannelMap().remove(userId);
    }

}
